import json
import math
import requests

# ========== 1. 定义本地计算器函数 ==========
def local_calculator(expression: str) -> dict:
    try:
        # 只允许 math 库里的安全函数
        allowed_names = {k: v for k, v in math.__dict__.items() if not k.startswith("__")}
        result = eval(expression, {"__builtins__": None}, allowed_names)
        return {"result": result}
    except Exception as e:
        return {"error": str(e)}

# ========== 2. 定义函数调用规范 ==========
functions = [
    {
        "name": "local_calculator",
        "description": "执行一个数学表达式的计算，例如 '2+3*4' 或 'sqrt(16)'",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "要计算的数学表达式，支持 +, -, *, /, sqrt 等"
                }
            },
            "required": ["expression"]
        }
    }
]

# ========== 3. 讯飞星火 API 配置 ==========
APIPassword = "Bearer QdoSJjUpKEotpPRZefVO:aaJeSlTMCGWZIBKQhkTc"
url = "https://spark-api-open.xf-yun.com/v2/chat/completions"

# 发送请求到讯飞星火 API
def send_request_to_xinghuo(messages, use_functions=False, stream=False):
    headers = {
        'Authorization': APIPassword,
        'content-type': "application/json"
    }
    
    body = {
        "model": "x1",
        "user": "user_id",
        "messages": messages,
        "stream": stream
    }
    
    # 如果需要使用函数调用，添加函数定义
    if use_functions:
        body["tools"] = [
            {
                "type": "function",
                "function": func
            } for func in functions
        ]
        body["tool_choice"] = "auto"
    
    try:
        response = requests.post(url=url, json=body, headers=headers, stream=stream)
        if response.status_code != 200:
            return f"API请求失败: {response.status_code}"
        
        # 非流式响应处理
        if not stream:
            data = response.json()
            return data
        
        # 流式响应处理 (这里简化处理，实际应用中可能需要更复杂的逻辑)
        full_response = ""
        for chunk in response.iter_lines():
            if chunk and b'[DONE]' not in chunk:
                if len(chunk) > 6:
                    data_org = chunk[6:]
                    try:
                        chunk_data = json.loads(data_org)
                        if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
                            choice = chunk_data['choices'][0]
                            if 'delta' in choice and 'content' in choice['delta']:
                                full_response += choice['delta']['content']
                    except json.JSONDecodeError:
                        pass
        return {"choices": [{"message": {"content": full_response}}]}
        
    except requests.exceptions.RequestException as e:
        return f"请求异常: {str(e)}"
    except Exception as e:
        return f"发生未知错误: {str(e)}"

# ========== 4. 用户输入 ==========
print("请输入您的问题或计算请求（例如：请帮我计算 sqrt(16)+3*5）：")
user_input = input("您的输入: ")

# ========== 5. 请求模型，获取函数调用建议 ==========
messages = [{"role": "user", "content": user_input}]
print("正在访问讯飞星火大模型，获取函数调用建议...")
response = send_request_to_xinghuo(messages, use_functions=True, stream=False)

# ========== 6. 处理响应结果 ==========
if isinstance(response, dict) and "choices" in response and len(response["choices"]) > 0:
    message = response["choices"][0]["message"]
    
    # 检查是否有工具调用
    if "tool_calls" in message and len(message["tool_calls"]) > 0:
        tool_call = message["tool_calls"][0]
        if tool_call["type"] == "function":
            fn_name = tool_call["function"]["name"]
            args = json.loads(tool_call["function"]["arguments"])
            
            if fn_name == "local_calculator":
                # 执行本地计算器函数
                result = local_calculator(**args)
                print("🧮 本地计算器执行结果:", result)
                
                # 将函数执行结果再交给模型解释
                messages.append(message)  # 添加模型提出的工具调用
                messages.append({
                    "role": "tool",
                    "name": fn_name,
                    "content": json.dumps(result)
                })  # 添加函数执行结果
                
                print("正在将函数执行结果提交给讯飞星火大模型进行解释...")
                final_response = send_request_to_xinghuo(messages, stream=True)
                
                if isinstance(final_response, dict) and "choices" in final_response:
                    print("🤖 模型最终回答:", final_response["choices"][0]["message"]["content"])
                else:
                    print("🤖 模型最终回答:", final_response)
    else:
        # 模型直接回复，没有调用函数
        print("模型没有选择调用函数，直接回复:", message["content"])
else:
    print("响应处理失败:", response)
